from ssummolib import load_index
from Bio import SeqIO
import os, re, sys, time
import CONFIG
import subprocess


top = CONFIG.arbDBdir

def parse_hmm():
        """This will parse the tabular result file from hmmsearch (must 
        specify --tblout temp_table.txt ) and return the Evalue and
        the score.
        """
#       x = re.compile('\s[0-9\.]+[e\-]*[0-9]*')
        x = re.compile('[a-zA-Z0-9\.\-]+\s+')
        score_list = []
        handle = file('temp_table.txt','r')
        one = 0
        for line in handle.xreadlines():
                if line.startswith('#'):
                        continue
                else:
                        if one == 1:
                                print 'more than one result in the table...'
                        string = line
                        one = 1
        handle.close()
        if one == 0:
                return False
        for result in x.finditer(string):
                result = result.group().strip()
                score_list.append(result)
        print "Taking {0} as Evalue and {1} as score".format(score_list[4], score_list[5])
        return score_list[4], score_list[5]

def SSUMMO(tdict):
        hmmsearch = os.path.join(CONFIG.hmmerdir, 'hmmsearch')
        ## At each path, choose which dir to go into, by performing
        ## hmmsearch on the fasta sequence against each next HMM.
        node = tdict
        top_len = len(CONFIG.arbDBdir)
        path = CONFIG.arbDBdir
        if 'accessions' in node:
                del( node['accessions'] )
        while len(node.keys()) > 0:
                temp_dict = {}
                for key in node.keys():
                        if key == 'accessions':
                                continue
                        temp_path = path + '/' + key
                        subprocess.call([hmmsearch,'--tblout','temp_table.txt','-o',os.devnull,'--noali',os.path.join(temp_path,key+'.hmm'),'temp.fas'],shell=False)
                        x = parse_hmm()
                        if x != False:
                                Evalue, score = x
                        else:
                                print 'no hmm results. Cant do this node: {0}'.format( os.path.join(path,key))
                                continue
                        temp_dict.update( { key : [float(Evalue), float(score)] } )
                best_key = None
                next_orgs = sorted( temp_dict.keys() )
                scores, Evalues = [], []
                for key in next_orgs:
                        Evalues.append(temp_dict[key][0])
                        scores.append(temp_dict[key][1])
                try:
                        best_Evalue = min(Evalues)
                        best_score = max(scores)
                except (UnboundLocalError, ValueError):
                        print "No scores found"
                        return path, 0, 0
                if scores.count(best_score) > 1 and Evalues.count(best_Evalue) > 1:
                        return path, 0, 0
                elif scores.count(best_score) == 1:
                        best_key = next_orgs[scores.index(best_score)]
                if best_key != None:
                        node = node[best_key]
                        path = path + '/' + best_key
                        print "The winner was %s with score: %s and Evalue: %s" % (path[top_len:], best_score, best_Evalue)
                else:
                        print "No result for sequence:-"
                        return None
        
        print "We got a winner!!! ", path[top_len:], "Score: ",best_score, "Evalue: ",best_Evalue
        return path, best_score, best_Evalue

if __name__ == '__main__':
        """Provide a sequence file as first command line arg and it shall be parsed one sequence
        at a time, and then SSUMMO shall be run on it...
        """
        try:
                seq_file_name = sys.argv[1]
        except IndexError:
                print "Wrong number of arguments"
                raise IndexError("Please provide a filename")
        out_name = seq_file_name[:seq_file_name.rfind('.')] + '_results.txt'
        results_file = file(out_name,'w')
        results_file.write('%s\t%s\t%s\t%s\t%s\t%s\n' % ('Name from fasta file', 'Genus guessed from SSUMMO','Seq. Length','Time','Score','Evalue'))
        seq_file = file(seq_file_name,'rU')
        t0 = time.time()
        seq_count = 0
        all_seqs = SeqIO.parse(seq_file,'fasta')
        dna_alphabet = re.compile('[ACGTUNSacgtuns]+')
        dna_res = 'ACGTUNS'
        for seq in all_seqs:
                org = seq.description
                org = org[org.rfind('[') + 1: org.rfind(']')]
                print "SSUMMOing %s" % org
                sequence = seq.seq.tostring()
                re_sult = dna_alphabet.search(sequence)
                result_len = len(re_sult.group())
                if re_sult == None or result_len != len(sequence):
                        print "Skipping %s 'cos one residue's not a nucleotide. Check it's first 80chars" % org
                        print sequence[:80]
                        if result_len > 5 and sequence[result_len] not in dna_res:
                                dna_res += sequence[result_len]
                        print "Offending seq (+4):",sequence[result_len:result_len+5]
                        continue
                else:
                        pass
                one_seq_file = file('temp.fas','w')
                one_seq_file.write(seq.format('fasta'))
                one_seq_file.close()
                try:
                        x = SSUMMO(load_index())
                        if x == None:
                                print '\tID: %s\n\tLength: %s' % (seq.description, len(sequence))
                                continue
                        else:
                                path_to_SSUMMO_org, score, Evalue = x
                except Exception, e:
                        results_file.close()
                        raise
                t = time.time() - t0
                print t
                results_file.write('%s\t%s\t%s\t%s\t%s\t%s\n' % (org,path_to_SSUMMO_org,len(seq),t,score, Evalue))
                seq_count += 1
        results_file.close()
        print "processed %s sequences in %s seconds" % (seq_count,t)
        print "DNA residues encountered: %s" % dna_res


